"""
WorldQuant 101 因子机器学习策略
使用线性回归模型预测收益率，基于预测值的 Z-score 生成交易信号
"""
import numpy as np
import pandas as pd
from sklearn.linear_model import LinearRegression
from sklearn.preprocessing import StandardScaler
from .wq101_zscore_strategy import WQ101ZscoreStrategy


class WQ101MLStrategy(WQ101ZscoreStrategy):
    strategy_name = "wq101_ml"
    
    def __init__(self, api, symbol_info, market_period):
        super().__init__(api, symbol_info, market_period)
        for symbol in symbol_info.keys():
            # 存储历史因子原始值序列
            self.all_info[symbol]["alpha001_series"] = []
            self.all_info[symbol]["alpha002_series"] = []
            self.all_info[symbol]["alpha003_series"] = []
            self.all_info[symbol]["alpha006_series"] = []
            self.all_info[symbol]["alpha011_series"] = []
            self.all_info[symbol]["alpha020_series"] = []
            self.all_info[symbol]["alpha023_series"] = []
            self.all_info[symbol]["alpha027_series"] = []
            self.all_info[symbol]["alpha030_series"] = []
            self.all_info[symbol]["alpha032_series"] = []
            self.all_info[symbol]["alpha041_series"] = []
            self.all_info[symbol]["alpha045_series"] = []
            self.all_info[symbol]["alpha053_series"] = []
            self.all_info[symbol]["alpha057_series"] = []
            self.all_info[symbol]["alpha065_series"] = []
            self.all_info[symbol]["alpha081_series"] = []
            self.all_info[symbol]["alpha084_series"] = []
            self.all_info[symbol]["alpha093_series"] = []
            self.all_info[symbol]["alpha094_series"] = []
            self.all_info[symbol]["alpha101_series"] = []
            
            # 存储收益率序列
            self.all_info[symbol]["return_series"] = []
            
            # 存储预测收益率序列（用于计算 Z-score）
            self.all_info[symbol]["predicted_return_series"] = []
            
            # 模型相关参数
            self.all_info[symbol]["lookback_period"] = symbol_info[symbol].get("lookback_period", 200)
            self.all_info[symbol]["train_period"] = symbol_info[symbol].get("train_period", 300)
            self.all_info[symbol]["zscore_window"] = symbol_info[symbol].get("zscore_window", 300)
            self.all_info[symbol]["max_series_length"] = symbol_info[symbol].get("max_series_length", 800)
            
            # 线性回归模型和标准化器
            self.all_info[symbol]["model"] = None
            self.all_info[symbol]["scaler"] = StandardScaler()
    
    def calculate_signal(self, symbol, current_pos, kline):
        """计算交易信号 - 基于预测收益率的 Z-score"""
        # 更新所有因子的原始值序列
        self._update_alpha_series(symbol, kline)
        
        # 更新收益率序列
        self._update_return_series(symbol, kline)
        
        # 获取参数
        lookback_period = self.all_info[symbol]["lookback_period"]
        train_period = self.all_info[symbol]["train_period"]
        
        # 检查是否有足够的数据进行训练
        min_data_length = lookback_period + train_period
        if len(self.all_info[symbol]["alpha001_series"]) < min_data_length:
            return 0, 0.0
        
        # 训练模型并预测
        predicted_return = self._train_and_predict(symbol)
        
        if predicted_return is None:
            return 0, 0.0
        
        # 保存预测收益率
        self.all_info[symbol]["predicted_return_series"].append(predicted_return)
        if len(self.all_info[symbol]["predicted_return_series"]) > self.all_info[symbol]["max_series_length"]:
            self.all_info[symbol]["predicted_return_series"].pop(0)
        
        # 计算预测收益率的 Z-score
        signal_score = self._calculate_prediction_zscore(symbol)
        
        # 获取 tanh 缩放系数
        tanh_scale = self.all_info[symbol].get("tanh_scale", 1.0)
        
        # 使用 tanh 将 Z-score 映射到 [-1, 1] 区间作为目标仓位
        target_position = np.tanh(signal_score * tanh_scale)
        
        # DEBUG: 打印前几次的预测
        if len(self.all_info[symbol]["predicted_return_series"]) <= 50:
            print(f"Bar {len(self.all_info[symbol]['predicted_return_series'])}: "
                  f"predicted_return={predicted_return:.6f}, zscore={signal_score:.3f}, "
                  f"target_pos={target_position:.3f}")
        
        return target_position, signal_score

    
    def _update_return_series(self, symbol, kline):
        """更新收益率序列"""
        close_prices = kline.close.values
        
        if len(close_prices) < 2:
            return
        
        # 计算收益率
        returns = np.diff(close_prices) / close_prices[:-1]
        
        # 更新收益率序列
        current_len = len(returns)
        stored_len = len(self.all_info[symbol]["return_series"])
        
        if stored_len < current_len:
            # 批量添加缺失的收益率
            self.all_info[symbol]["return_series"].extend(returns[stored_len:].tolist())
            
            # 限制序列长度
            max_length = self.all_info[symbol]["max_series_length"]
            if len(self.all_info[symbol]["return_series"]) > max_length:
                self.all_info[symbol]["return_series"] = self.all_info[symbol]["return_series"][-max_length:]
    
    def _train_and_predict(self, symbol):
        """训练模型并预测下一期收益率"""
        lookback_period = self.all_info[symbol]["lookback_period"]
        train_period = self.all_info[symbol]["train_period"]
        
        # 获取所有因子序列
        series_dict = {
            'alpha001': self.all_info[symbol]["alpha001_series"],
            'alpha002': self.all_info[symbol]["alpha002_series"],
            'alpha003': self.all_info[symbol]["alpha003_series"],
            'alpha006': self.all_info[symbol]["alpha006_series"],
            'alpha011': self.all_info[symbol]["alpha011_series"],
            'alpha020': self.all_info[symbol]["alpha020_series"],
            'alpha023': self.all_info[symbol]["alpha023_series"],
            'alpha027': self.all_info[symbol]["alpha027_series"],
            'alpha030': self.all_info[symbol]["alpha030_series"],
            'alpha032': self.all_info[symbol]["alpha032_series"],
            'alpha041': self.all_info[symbol]["alpha041_series"],
            'alpha045': self.all_info[symbol]["alpha045_series"],
            'alpha053': self.all_info[symbol]["alpha053_series"],
            'alpha057': self.all_info[symbol]["alpha057_series"],
            'alpha065': self.all_info[symbol]["alpha065_series"],
            'alpha081': self.all_info[symbol]["alpha081_series"],
            'alpha084': self.all_info[symbol]["alpha084_series"],
            'alpha093': self.all_info[symbol]["alpha093_series"],
            'alpha094': self.all_info[symbol]["alpha094_series"],
            'alpha101': self.all_info[symbol]["alpha101_series"],
        }
        
        return_series = self.all_info[symbol]["return_series"]
        
        # 检查数据长度 - 确保所有序列长度一致
        min_series_len = min(len(s) for s in series_dict.values())
        
        # 确保有足够的数据进行训练
        min_data_required = lookback_period + train_period
        if min_series_len < min_data_required or len(return_series) < min_data_required:
            return None
        
        # 使用较短序列的长度作为基准
        available_len = min(min_series_len, len(return_series))
        
        # 构建训练数据
        X_train_list = []
        y_train_list = []
        
        # 滚动窗口构建训练样本
        start_idx = available_len - train_period - lookback_period
        end_idx = available_len - lookback_period
        
        for i in range(start_idx, end_idx):
            # 检查索引是否有效
            if i < 0 or i + lookback_period >= len(return_series):
                continue
            
            # 获取前 lookback_period 期的因子值作为特征
            features = []
            for key in sorted(series_dict.keys()):
                if i + lookback_period > len(series_dict[key]):
                    continue
                features.extend(series_dict[key][i:i+lookback_period])
            
            # 检查特征长度是否正确
            expected_feature_len = lookback_period * len(series_dict)
            if len(features) != expected_feature_len:
                continue
            
            X_train_list.append(features)
            # 目标是下一期的收益率
            y_train_list.append(return_series[i + lookback_period])
        
        # 检查是否有足够的训练样本
        if len(X_train_list) < 10:
            print(f"警告：训练样本不足 ({len(X_train_list)} 个)，跳过训练")
            return None
        
        X_train = np.array(X_train_list)
        y_train = np.array(y_train_list)
        
        # 检查是否有 NaN 或 Inf
        if np.any(np.isnan(X_train)) or np.any(np.isinf(X_train)):
            print(f"警告：训练数据包含 NaN 或 Inf，跳过训练")
            return None
        
        if np.any(np.isnan(y_train)) or np.any(np.isinf(y_train)):
            print(f"警告：目标数据包含 NaN 或 Inf，跳过训练")
            return None
        
        # 标准化特征
        X_train_scaled = self.all_info[symbol]["scaler"].fit_transform(X_train)
        
        # 训练线性回归模型
        model = LinearRegression()
        model.fit(X_train_scaled, y_train)
        self.all_info[symbol]["model"] = model
        
        # 预测：使用最新的 lookback_period 期因子值
        # 确保有足够的数据
        if any(len(series_dict[key]) < lookback_period for key in series_dict.keys()):
            return None
        
        features_current = []
        for key in sorted(series_dict.keys()):
            features_current.extend(series_dict[key][-lookback_period:])
        
        # 检查特征是否有 NaN 或 Inf
        if np.any(np.isnan(features_current)) or np.any(np.isinf(features_current)):
            print(f"警告：预测特征包含 NaN 或 Inf")
            return None
        
        X_current = np.array([features_current])
        X_current_scaled = self.all_info[symbol]["scaler"].transform(X_current)
        
        predicted_return = model.predict(X_current_scaled)[0]
        
        return predicted_return
    
    def _calculate_prediction_zscore(self, symbol):
        """计算预测收益率的 Z-score"""
        predicted_series = self.all_info[symbol]["predicted_return_series"]
        
        if len(predicted_series) < 2:
            return 0.0
        
        zscore_window = self.all_info[symbol]["zscore_window"]
        
        # 使用最近 zscore_window 个预测值计算均值和标准差
        recent_predictions = predicted_series[-zscore_window:] if len(predicted_series) >= zscore_window else predicted_series
        
        mean_pred = np.mean(recent_predictions)
        std_pred = np.std(recent_predictions)
        
        if std_pred == 0:
            return 0.0
        
        # 当前预测值的 Z-score
        current_prediction = predicted_series[-1]
        zscore = (current_prediction - mean_pred) / std_pred
        
        # 限制在 [-5, 5] 范围内
        return np.clip(zscore, -5, 5)


